## plots description of MC data under various scenarios
library(mvtnorm)
library(lattice)
library(mgcv)

VC <- diag(3)


n.vec <- c(250, 500, 1000)
alpha.vec <- rbind(c(.1,.1,.05),
                   c(1, 1, .5),
                   c(1.5,1.5,.75))
beta.vec <- rbind(c(1,1,2,0,0,5),
                  c(1,1,2,2,2,1))
                   
                   

set.seed(1234)


n <- 500000
conf <- c("Confounding = Low", "Confounding = Moderate", "Confounding = Severe")
nonl <- c("Outcome Mean = Linear", "Outcome Mean = Nonlinear")
pscore.data <- NULL
for (alpha.ind in 1:3){
    
  alpha <- alpha.vec[alpha.ind,]
    
  Z <- rmvnorm(n, mean=rep(0, 3), sigma=VC)
  
  eta <- Z[,1]*alpha[1] + Z[,2]*alpha[2] +
    Z[,1]*Z[,2]*alpha[3]
  X <- rbinom(n, 1, pnorm(eta))
  p <- pnorm(eta)
  
  holder <- data.frame(p=p, x=X, Confounding=conf[alpha.ind],
                       Assignment="Assignment Model = True")
  pscore.data <- rbind(pscore.data, holder)
  
}
n <- 50000
for (alpha.ind in 1:3){
    
  alpha <- alpha.vec[alpha.ind,]
    
  Z <- rmvnorm(n, mean=rep(0, 3), sigma=VC)
  
  eta <- Z[,1]*alpha[1] + Z[,2]*alpha[2] +
    Z[,1]*Z[,2]*alpha[3]
  X <- rbinom(n, 1, pnorm(eta))

  gam.out <- gam(X~s(Z[,2]), family=binomial(probit))
  p <- predict(gam.out, type="response")
 
  
  holder <- data.frame(p=p, x=X, Confounding=conf[alpha.ind],
                       Assignment="Assignment Model = Minimal")
  pscore.data <- rbind(pscore.data, holder)
  
}

trellis.device("pdf", col=FALSE)
pdf("ConfoundingSummaryMC.pdf", height=8, width=12, bg="white")
supline <- trellis.par.get("superpose.line")
supline$lty <- c(1,2)
trellis.par.set("superpose.line", supline)
print(densityplot(~p|Confounding*Assignment,
                  groups=x, data=pscore.data, layout=c(3,2),
                  plot.points=FALSE, lty=c(1,2),
                  key=list(text=list(c("Control", "Treated")),
                    lines=Rows(trellis.par.get("superpose.line"), c(1,2))
                    ),
                  xlab="Probability of Treatment"))
dev.off()




plot.data <- NULL
for (n in n.vec){
  for (alpha.ind in 1:3){
    for (beta.ind in 1:2){
      
      alpha <- alpha.vec[alpha.ind,]
      beta <- beta.vec[beta.ind,]
      
      Z <- rmvnorm(n, mean=rep(0, 3), sigma=VC)
      
      eta <- Z[,1]*alpha[1] + Z[,2]*alpha[2] +
        Z[,1]*Z[,2]*alpha[3]
      X <- rbinom(n, 1, pnorm(eta))
      Y <- Z[,2]*beta[1] + Z[,3]*beta[2] + Z[,2]*X*beta[3] +
        (Z[,2]^2)*X*beta[4] + (Z[,3]^2)*X*beta[5] + X*beta[6] + rnorm(n,s=1)
      
      holder <- data.frame(z=Z[,2], x=X, y=Y, n=n,
                           Confounding=conf[alpha.ind],
                           Nonlinearity=nonl[beta.ind])
     
      plot.data <- rbind(plot.data, holder)
      
    }
  }
  if (n == 250){
    pdf("MCExample250.pdf", height=4, width=8, bg="white")
    print(xyplot(y~z|Confounding*Nonlinearity, data=plot.data,
                 xlab="z2",
                 groups=x, subset=plot.data$n==250, pch=1, cex=.2))
    dev.off()
  }
  if (n == 500){
    pdf("MCExample500.pdf", height=4, width=8, bg="white")
    print(xyplot(y~z|Confounding*Nonlinearity, data=plot.data,
                 xlab="z2",
                 groups=x, subset=plot.data$n==500, pch=1, cex=.2))
    dev.off()
  }
  if (n == 1000){
    pdf("MCExample1000.pdf", height=4, width=8, bg="white", version="1.4")
    print(xyplot(y~z|Confounding*Nonlinearity, data=plot.data,
                 xlab=expression(z[2]), col=c(rgb(0,0,0,.5),
                     rgb(0.65, 0.65, 0.65,.5)),
                 groups=x, subset=plot.data$n==1000, pch=1, cex=.2))
    dev.off()
  }
  
}



